# 计算指数与个股的相关度
from datacache import get_index_data, get_stock_data

def calc_index_stock_corr(stock_code, index_code):
    start_date = "2020-05-05"
    end_date = "2025-02-01"
    formatted_start_date = start_date.replace("-", "")
    formatted_end_date = end_date.replace("-", "")
    print("开始计算")
    index_data = get_index_data(index_code, start_date, end_date)
    print(index_data)
    stock_data = get_stock_data(stock_code, formatted_start_date, formatted_end_date)
    print(stock_data)
    if len(index_data) != len(stock_data):
        print("数据长度不一致，无法计算相关度")
        return
    correlation = index_data['收盘'].corr(stock_data['close'])
    print(f"股票代码{stock_code}与大盘{index_code}的相关度为：{correlation}")
    return correlation

if __name__ == '__main__':
    stock_code_base = "sz000001"
    index_code_base = "000001"
    calc_index_stock_corr(stock_code=stock_code_base, index_code=index_code_base)